#define _CRT_SECURE_NO_WARNINGS

#include"sm2.h"

#include <string.h>
#include <malloc.h>
#include <stdlib.h>

#include "pointoperation.h"
#include "kdf.h"
#include "sm3.h"

/**
 * 初始化SM2椭圆曲线参数.
 * 
 * \param ecc_param
 */
void sm2_param_init(param sm2_param) {
    // 分配空间
    mpz_inits(sm2_param->p, sm2_param->a, sm2_param->b, sm2_param->n, NULL);

    // 初始化
    // 采用16进制字符串
    int base = 16;

    // p = BDB6F4FE 3E8B1D9E 0DA8C0D4 6F4C318C EFE4AFE3 B6B8551F
    mpz_set_str(sm2_param->p, "BDB6F4FE3E8B1D9E0DA8C0D46F4C318CEFE4AFE3B6B8551F", base);

    // a = BB8E5E8F BC115E13 9FE6A814 FE48AAA6 F0ADA1AA 5DF91985
    mpz_set_str(sm2_param->a, "BB8E5E8FBC115E139FE6A814FE48AAA6F0ADA1AA5DF91985", base);

    // b = 1854BEBD C31B21B7 AEFC80AB 0ECD10D5 B1B3308E 6DBF11C1
    mpz_set_str(sm2_param->b, "1854BEBDC31B21B7AEFC80AB0ECD10D5B1B3308E6DBF11C1", base);

    // n = BDB6F4FE 3E8B1D9E 0DA8C0D4 0FC96219 5DFAE76F 56564677
    mpz_set_str(sm2_param->n, "BDB6F4FE3E8B1D9E0DA8C0D40FC962195DFAE76F56564677", base);
}

/**
 * 释放SM2椭圆曲线参数的空间.
 *
 * \param sm2_param
 */
void sm2_param_clear(param sm2_param) {
    // 分配空间
    mpz_clears(sm2_param->p, sm2_param->a, sm2_param->b, sm2_param->n, NULL);
}

/**
 * 根据私钥private_key 和 基点G生成公钥.
 * 
 * \param public_key
 * \param private_key
 * \param sm2_param
 * \param base_point
 */
void public_key_gen(point public_key, mpz_t private_key, param sm2_param, point  base_point) {
    point_mul(public_key, base_point, private_key, sm2_param->a, sm2_param->p);
}


/**
 * SM2加密.
 * 
 * \param ct 密文
 * \param ct_len 密文长度
 * \param sm2_param SM2的参数
 * \param public_key 公钥
 * \param base_point 基点G
 * \param msg 消息（明文）
 * \param msg_len 明文（明文）的长度
 */
void sm2_encrypt(uint8_t* ct, size_t ct_len, param sm2_param, point public_key, point  base_point, uint8_t* msg, size_t msg_len) {
    // 初始化
   // 字符串使用16进制
    size_t base = 16;
    // 坐标的长度 48个16进制数
    // size_t position_len = 48;
    size_t position_len = 24;
    // mpz转换成字符串之后的长度，符号+结尾 （+2字节）。但是这里不考虑符号，所以只+1
    size_t mpz_str_len = position_len + 1;

    size_t ct_index = 0;

    // 随机数
    // 384F3035 3073AEEC E7A16543 30A96204 D37982A3 E15B2CB5
    mpz_t rand_k;
    mpz_init_set_str(rand_k, "384F30353073AEECE7A1654330A96204D37982A3E15B2CB5", base);

    // 1.生成密文C1
    // 1.1.初始化多倍点 k[G]
    point k_g;
    point_init(k_g);

    // 1.2.计算 (x1, y2) = k[G] =rand_k[G]
    point_mul(k_g, base_point, rand_k, sm2_param->a, sm2_param->p);

    // 1.3.计算 C1 = 04 || x1 || y1
    ct[ct_index] = 0x04;
    ct_index = ct_index + 1;

    // mpz_to_fixed_length_str(ct + 2, mpz_str_len, k_g->x, base);
    mpz_to_byte_array(ct + ct_index, position_len, k_g->x);
    ct_index = ct_index + position_len;

    // mpz_to_fixed_length_str(ct + 2 + position_len, mpz_str_len, k_g->y, base);
    mpz_to_byte_array(ct + ct_index, position_len, k_g->y);
    ct_index = ct_index + position_len;

    // print_chars_by_hex("打印ct1", ct + 1, 2 + position_len * 2);

    // 2.生成密文C2
    // 2.1.计算多倍点 (x2, y2) = k[Pb] = rand_k[Pb]
    point k_public_key;
    point_init(k_public_key);
    point_mul(k_public_key, public_key, rand_k, sm2_param->a, sm2_param->p);

    size_t x2_y2_len = 2 * position_len;
    uint8_t* x2_y2 = (uint8_t*)malloc(sizeof(uint8_t) * x2_y2_len);

    // mpz_to_fixed_length_str(x2_y2, mpz_str_len, k_public_key->x, base);
    mpz_to_byte_array(x2_y2, position_len, k_public_key->x);
    // mpz_to_fixed_length_str(x2_y2 + position_len, mpz_str_len, k_public_key->y, base);
    mpz_to_byte_array(x2_y2 + position_len, position_len, k_public_key->y);
    // print_chars_by_hex("x2_y2", x2_y2, x2_y2_len);

    // 2.2.使用密钥派生函数，生成 t = KDF(x2 || y2, len) ,其中 len = klen/8，klen是8的倍数
    uint8_t* ct2 = (uint8_t*)malloc(sizeof(uint8_t) * msg_len);

    kdf(ct2, msg_len, x2_y2, x2_y2_len);
    // print_chars_by_hex("r = KDF(x2 ‖ y2, klen)", ct2, msg_len);

    // 2.3.计算 C2 = M ^ t
    for (size_t i = 0; i < msg_len; i++)
    {
        ct2[i] = msg[i] ^ ct2[i];
    }
    // print_chars_by_hex("ct2", ct2, msg_len);

    // 3.生成密文C3
    // C3 = Hash(x2 || M || y2)
    size_t ct3_temp_len = position_len * 2 + msg_len;
    uint8_t* ct3_temp = (uint8_t*)malloc(sizeof(uint8_t) * ct3_temp_len);
    uint8_t ct3[32];

    memcpy(ct3_temp, x2_y2, position_len);
    memcpy(ct3_temp + position_len, msg, msg_len);
    memcpy(ct3_temp + position_len + msg_len, x2_y2 + position_len, position_len);

    // print_chars_by_hex("ct3_temp", ct3_temp, ct3_temp_len);

    sm3_context ctx;

    sm3_init(&ctx);

    sm3_update(&ctx, ct3_temp, ct3_temp_len);

    sm3_finish(ct3, &ctx);

    // 4.输出结果
    // ct = C1 || C3 || C2
    memcpy(ct + ct_index, ct3, 32);
    ct_index = ct_index + 32;

    memcpy(ct + ct_index, ct2, msg_len);
    ct_index = ct_index + msg_len;

    // 5.释放空间
    free(ct3_temp);

    free(x2_y2);

    free(ct2);
}


/**
 * SM2解密.
 * 
 * \param plaintext 明文
 * \param plaintext_len 明文长度
 * \param sm2_param SM2的参数
 * \param private_key 私钥
 * \param ct 密文
 * \param ct_len 密文长度
 */
void sm2_decrypt(uint8_t* plaintext, size_t plaintext_len, param sm2_param, mpz_t private_key, uint8_t* ct, size_t ct_len)
{
    // TODO 1.验证是否满足方程上的点，不满足退出
    // TODO 2.验证点 S=H[C1],若 S 是无穷远点 ,则 报错并退出;
    // 3.计算 db[C1]=private_key[C1]
    point k_g, private_key_c1;
    point_inits(k_g, private_key_c1, NULL);

    // 恢复k_g
    // 坐标的字节数
    size_t position_len = 24;
    // ct的索引
    size_t ct_index = 1;

    mpz_import(k_g->x, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, ct + ct_index);
    ct_index = ct_index + position_len;
    // gmp_printf("kg=%ZX\n", k_g->x);

    mpz_import(k_g->y, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, ct + ct_index);
    ct_index = ct_index + position_len;

    // 计算 db[C1]=private_key[C1]
    point_mul(private_key_c1, k_g, private_key, sm2_param->a, sm2_param->p);

    // 4.计算 t=KDF(x2||y2,klen)，若t全为0，报错退出
    size_t x2_y2_len = 2 * position_len;
    uint8_t* x2_y2 = (uint8_t*)malloc(sizeof(uint8_t) * x2_y2_len);

    mpz_to_byte_array(x2_y2, position_len, private_key_c1->x);
    mpz_to_byte_array(x2_y2 + position_len, position_len, private_key_c1->y);
    // print_chars_by_hex("x2_y2", x2_y2, x2_y2_len);

    kdf(plaintext, plaintext_len, x2_y2, x2_y2_len);
    // print_chars_by_hex("KDF plaintext", plaintext, plaintext_len);

    // 5.计算出 M'=C2 ^ t
    ct_index = ct_index + 32;
    for (size_t i = 0; i < plaintext_len; i++)
    {
        plaintext[i] = plaintext[i] ^ ct[i + ct_index];
    }

    // TODO 6.计算出 u = hash(x2 || M1 || y2)，若 u != C3 退出

    // 释放空间
    free(x2_y2);
}

/**
 * 将mpz转换为固定位数的字符串.
 * 
 * \param num mpz格式的数
 * \param length 输出的字符串的长度
 * \param base 基数
 * \return 
 */
void mpz_to_fixed_length_str(char* padded_str,size_t padded_str_len ,mpz_t num, int base) {
    // 先将 mpz 的数据转换为 str
    char* raw_str = malloc(sizeof(char) * padded_str_len);

    mpz_get_str(raw_str, base, num);


    int raw_len = strlen(raw_str);

    // 计算需要填充的零的数量
    // -1 是因为有一个长度给结束符号的
    int padding_len = padded_str_len - 1 - raw_len;
    if (padding_len <= 0) {
        // 实际长度大于等于所需长度，无需填充
        strcpy(padded_str, raw_str);

        // 释放原始字符串占用的内存
        free(raw_str);

        return ;
    }

    // 填充零
    memset(padded_str, '0', padding_len);

    // 复制原始数字字符串到填充后的字符串中
    strcpy(padded_str + padding_len, raw_str);

    // 释放原始字符串占用的内存
    free(raw_str);
}


/**
 * 用于将 mpz_t 类型转换为字节数组.
 * 
 * \param byte_array 输出数组
 * \param size 输出数组的长度（size必须大于等于mpz值占的字节数，不够的部分填充0）
 * \param value mpz值
 */
void mpz_to_byte_array(unsigned char* byte_array, size_t size, mpz_t value)
{
    size_t n, limbs;

    // 获取所需的字节数量
    n = mpz_sizeinbase(value, 256);

    // 获取 limb 数量
    limbs = (n + (sizeof(mp_limb_t) * 8) - 1) / (sizeof(mp_limb_t) * 8);

    // 导出到字节数组
    // 输出数组、转为256进制之后的长度，1表示不逆序，sizeof(mp_limb_t) 一个limb的大小，1小端，钉子位空间够默认为0，输出结果value

    // 计算填充位置
    size_t padding_len = size - n;
    if (padding_len > 0)
    {
        mpz_export(byte_array + padding_len, &n, 1, sizeof(mp_limb_t), 1, 0, value);
    }
    else
    {
        mpz_export(byte_array, &n, 1, sizeof(mp_limb_t), 1, 0, value);
    }
}
